###############################################################################
#
# Copyright (c) 2011 Rainmaker Entertainment
# All Rights Reserved.
#
# This file contains unpublished confidential and proprietary
# information of Rainmaker Entertainment.  The contents of this file
# may not be copied or duplicated, in whole or in part, by any
# means, electronic or hardcopy, without the express prior
# written permission of Rainmaker Entertainment.
#
#    $HeadURL: http://pt/src/trunk/maya/rnkAssetChecks/geo.py $
#    $Revision: 68280 $
#    $Author: wmakari $
#    $Date: 2013-09-30 15:38:06 -0700 (Mon, 30 Sep 2013) $
#
###############################################################################

###############################################################################
## INFOS
###############################################################################

# Geo checks


###############################################################################
## MODULES
###############################################################################

import traceback
import re
import math

try :
    import maya.cmds as cmds
    import maya.mel as mel
    import maya.OpenMaya as OpenMaya
except :
    traceback.print_exc()
    print "Can't load Maya modules."
    
    
import rnkAssetChecks.core as rnkAssetChecksCore
reload(rnkAssetChecksCore)
        


###############################################################################
## FUNCTIONS
###############################################################################  


def checkEmptyShapes(assetRoot):
    """ 
    @brief Check if there is any empty mesh shapes. Shave nodes are excluded.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    emptyShapesList = []
    
    getShapes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    if not len(getShapes) == 0 :
        for shape in getShapes :
            if 'shaveDisplayShape' not in shape and cmds.polyEvaluate(shape, v=True) == 0 :
                emptyShapesList.append(shape)


    return emptyShapesList
    


def checkUV(assetRoot) :
    """ 
    @brief Check to ensure that all meshes have UVs.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    result = []
    getShapes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    for mesh in getShapes :
        if cmds.polyEvaluate(mesh, uv=True) == 0 :
            result.append(mesh)
            
    return result



def checkNamesEndWithNumber(assetRoot):
    """ 
    @brief Check to ensure that mesh parent transform names do not end with a number.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """
    
    results = list()

    # get node to search
    nodes = cmds.listRelatives(assetRoot, ad=True)
    if not nodes:
        return list()

    # check node names
    for n in nodes:
        if re.match(r'(.*)([0-9]+)$', n):
            results.append(n)

    return results



def checkPivot(assetRoot):
    """ 
    @brief Check to ensure that all transforms have (0,0,0) rotation/scale pivots.
           Excluded from check: nodes ending with "_EX" and "follicle" type nodes.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    def valid(node):
        """
        Check if node's pivot should be zero,
        some cases like hair nodes, it doesn't matter.
        """
        if rnkAssetChecksCore.detectEyeBall(node):
            return True

        if node.endswith("_EX"):
            return False

        # Check by shapes
        shapes = cmds.listRelatives( node, shapes=True )

        if not shapes:
            return True

        # if it's nurbsCurve, check the pivot only when it's not driven by follicles
        nrbsCrv = cmds.listRelatives(node, shapes=True, type='nurbsCurve')
        if nrbsCrv :
            if cmds.listConnections("%s.tx" % node, s=True, d=False):
                x = cmds.listConnections(node, s=True, d=False)[0]
                if cmds.listRelatives(x, shapes=True, type="follicle"):
                    return False

            ## Check if its shape is not connected to an rnkFur or follicule
            if cmds.pluginInfo("rnkFur", q=True, loaded=True):
                for eachShape in nrbsCrv :
                    y = cmds.listConnections(eachShape, s=True, d=True)
                    if y :
                        for a in y :
                            if cmds.nodeType(a) == 'follicle'  or cmds.nodeType(a) == 'rnkFur' :
                                return False
                            else :
                                if cmds.listRelatives(a, shapes=True, type="follicle") :
                                    return False
                                elif cmds.listRelatives(a, shapes=True, type="rnkFur"):
                                    return False

        for s in shapes:
            if cmds.nodeType(s) in ["shape", "mesh", "nurbsSurface", 'nurbsCurve']:
                return True

            elif cmds.pluginInfo('rnkBeadDistributorNode', q=True, loaded=True) and cmds.nodeType(s) == "rnkBeadDistributor":
                return True


        return False



    badList = []
    transforms = cmds.listRelatives(assetRoot, ad=True, type="transform", f=True)
    
    if not len(transforms) == 0 :
        for trn in transforms:
            # this will avoid the eyebrows and eyelashes too...
        #if 'eye' not in trn:
            if not rnkAssetChecksCore.detectEyeBall(trn) and not cmds.nodeType(trn) == 'place3dTexture' :
    
                for pos in ['positionA', 'positionB'] :
                    if trn.endswith(pos):
                        print 'Check local pivot here for: %s' % trn
                        
                        
                ##--
                ##--WORLD SPACE
                ##--
                rotatePivot = cmds.xform( trn, q=True, ws=True, rotatePivot=True )
                if [ v for v in rotatePivot if v > 0.00001 or v <- 0.00001 ]:
                    if valid(trn) is True :
                        badList.append( trn )
                        continue
        
                scalePivot = cmds.xform( trn, q=True, ws=True, scalePivot=True )
                if [ v for v in scalePivot if v>0.00001 or v<-0.00001 ]:
                    if valid(trn) is True :
                        badList.append( trn )
                        continue
        
                ##--
                ##--LOCAL SPACE
                ##--
                rotatePivot = cmds.xform( trn, q=True, rotatePivot=True )
                if [ v for v in rotatePivot if v > 0.00001 or v <- 0.00001 ]:
                    if valid(trn) is True :
                        badList.append( trn )
                        continue
        
                scalePivot = cmds.xform( trn, q=True, scalePivot=True )
                if [ v for v in scalePivot if v>0.00001 or v<-0.00001 ]:
                    if valid(trn) is True :
                        badList.append( trn )
                        continue


    return badList



def isMeshTweaked(assetRoot):
    """ 
    @brief Check to ensure that all mesh vertices have (0,0,0) local transform.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    issues = []

    # get node to search
    nodes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    if not len(nodes) == 0 :
        for mesh in nodes:
            if not cmds.getAttr('%s.pnts' % mesh, size=True):
                continue

            meshMObj = rnkAssetChecksCore.asMObject(mesh)
    
            fn = OpenMaya.MFnDagNode(meshMObj)
            arrayPlug = fn.findPlug("pnts")
    
            clean = True
            for index in range(0, arrayPlug.numElements()):
                plug = arrayPlug.elementByPhysicalIndex(index)
                sum = 0.0
                for childIndex in range(0, 3):
                    sum = sum + math.fabs(plug.child(childIndex).asDouble())
    
                if sum > 1.0e-3:
                    clean = False
                    break
    
            if not clean:
                issues.append(mesh)


    return issues



def checkQuadPolygons(assetRoot):
    """ 
    @brief Check to ensure that all polygons have only 3 or 4 edges.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    issues = []

    meshes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    
    if meshes:
        cmds.select(meshes)
        cmd ='''polyCleanupArgList 3 { "1","2","1","0","1","0","0","0","0","1e-05","0","1e-05","0","1e-05","0","-1","0" };'''
        mel.eval(cmd)
        issues = cmds.ls(sl=True)
        cmds.select( cl=True )

    return issues


def checkFacesWithHoles(assetRoot):
    """ 
    @brief Check to ensure that no mesh contains holes.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    issues = []

    shapes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    if shapes:
        cmds.select(shapes)

        cmd ='''
        polyCleanupArgList 3 { "1","2","1","0","0","0","1","0","0","1e-05","0","0.001","0","1e-05","0","-1","0" };
        '''
        mel.eval(cmd)
        issues = cmds.ls(sl=True)
        cmds.select( cl=True )

    return issues


def checkZeroLengthEdges(assetRoot):
    """ 
    @brief Check to ensure that no mesh edges have zero length.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    issues = []

    shapes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    if shapes:
        cmds.select(shapes)
        cmd ='''
        polyCleanupArgList 3 { "1","2","1","0","0","0","0","0","0","1e-05","1","1e-05","0","1e-05","0","-1","0" };
        '''
        mel.eval(cmd)
        issues = cmds.ls(sl=True)
        cmds.select( cl=True )

    return issues


def checkNonmanifoldGeo(assetRoot):
    """ 
    @brief Check to ensure that all meshes are 2-manifold geometry.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    issues = []

    shapes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    if shapes:
        cmds.select(shapes)

        cmd ='''
        polyCleanupArgList 3 { "1","2","1","0","0","0","0","0","0","1e-05","0","1e-05","0","1e-05","0","1","0" };
        '''
        mel.eval(cmd)
        issues = cmds.ls(sl=True)
        cmds.select( cl=True )

    return issues



def checkTrianglesSharingVtx(assetRoot):
    """ 
    @brief Check to see if any triangles share one or more vertex.
    @return : List
    
    @param assetRoot - String - Asset root node.
    """

    if not cmds.pluginInfo("Mayatomr", q=True, l=True):
        cmds.loadPlugin("Mayatomr")

    issues = []

    shapes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    try :
        if shapes:
            for shape in shapes :


                # don't proceed if the parent transform has attribute __GrowthMesh
                transform = cmds.listRelatives(shape, p=True)[0]
                if cmds.attributeQuery("__GrowthMesh", n=transform, ex=True):
                    continue

                # skip the check if it's not connected to mentalraySubdivApprox
                if cmds.listConnections(shape, type='mentalraySubdivApprox'):
                    continue

                if cmds.listConnections(transform, type='mentalraySubdivApprox'):
                    continue


                cmds.select(shape)
                
                cmds.polySelectConstraint(mode=3, type=8, size=1)
                
                selFaces = cmds.ls(sl=True, flatten=True)
                if not len(selFaces) == 0 :
                    vtxDict = {}
                    for eachFace in selFaces :
                        selVtx = cmds.polyListComponentConversion(eachFace, toVertex=True)
                        selVtx = rnkAssetChecksCore.convertIdRangeToIdList(selVtx, convertFromSel=True, selType='vertex')
                        
                        for eachVtx in selVtx :
                            if vtxDict.has_key(eachVtx) is False :
                                vtxDict[eachVtx] = set()
                                
                            vtxDict[eachVtx].add(eachFace)
                            
                    for eachKey in vtxDict.keys() :
                        facesList = list(vtxDict[eachKey])
                        if len(facesList) > 1 :
                            for eachFace in facesList :
                                if not eachFace in issues :
                                    issues.append(eachFace)
                                    
        ## Turn constraint Off
        cmds.select(clear=True)
        cmds.polySelectConstraint(mode=0)
                 
        return issues          

    except :
        ## Turn constraint Off
        cmds.select(clear=True)
        cmds.polySelectConstraint(mode=0)
    
        print traceback.print_exc()
        
        
def useOldStyleSubdivs(assetRoot):
    """
    Adding miExportCCMesh to all mentalraySubdivApprox nodes. Affects the entire scene!
    """
    if not cmds.pluginInfo("Mayatomr", q=True, l=True):
        cmds.loadPlugin("Mayatomr")

    approx_nodes = cmds.ls(long=True, type=['mentalraySubdivApprox'])

    for approx_node in approx_nodes:

        if not cmds.attributeQuery('miExportCCMesh', node = approx_node, exists = True):
            cmds.addAttr(approx_node, ln='miExportCCMesh', at='bool')
            cmds.setAttr(approx_node + '.miExportCCMesh', False)
            
            
def checkOverlappingFaces(assetRoot):
    """
    Check if a geometry have been duplicated and then merged into one shape.
    """
    
    badGeos = []
    
    shapes = rnkAssetChecksCore.listShapes(assetRoot, types=["mesh"])
    trfsList = []
    if not len(shapes) == 0 :
        for shape in shapes :
            trfs = cmds.listRelatives(shape, parent=True)[0]
            if not trfs in trfsList :
                trfsList.append(trfs)
                
    if not len(trfsList) == 0 :
        for geometryName in trfsList :

            if overlappingFaces(geometryName) is True :
                badGeos.append(geometryName)

    
    return badGeos
    
    
def overlappingFaces(geometryName):
    
    ## Add the geometry in the selection
    selection = OpenMaya.MSelectionList()
    selection.add(geometryName)
    iterSel = OpenMaya.MItSelectionList(selection, OpenMaya.MFn.kMesh)
    
    ## Get DAG path
    dagPath = OpenMaya.MDagPath()
    iterSel.getDagPath(dagPath)      
    
    ## Get vertex list
    mitMesh = OpenMaya.MItMeshPolygon(dagPath)
    
    facePosDict = {}
    
    while not mitMesh.isDone():
    
            faceIdx = mitMesh.index()
            facePos = mitMesh.center(OpenMaya.MSpace.kWorld)
            faceNormal = OpenMaya.MVector()
            mitMesh.getNormal(faceNormal)
            
            facePosTpl = (facePos.x, facePos.y, facePos.z, facePos.w)
            if facePosTpl in facePosDict.keys() :
                if (faceNormal.x, faceNormal.y, faceNormal.z) in facePosDict[facePosTpl] :
                    return True
                    break
            else :
                if not facePosTpl in facePosDict :
                    facePosDict[facePosTpl] = [ (faceNormal.x, faceNormal.y, faceNormal.z) ]
                else :
                    list = facePosDict[facePosTpl]
                    list.append( (faceNormal.x, faceNormal.y, faceNormal.z) )
                    facePosDict[facePosTpl] = list
    
    
            mitMesh.next()
            
    return False